package com.zrc.japktool.apk.merge;

import com.zrc.japktool.apk.merge.dex.ModifiyDexClassID1;
import com.zrc.japktool.util.Zip4jUtil;
import net.lingala.zip4j.ZipFile;
import net.lingala.zip4j.model.FileHeader;
import org.jf.dexlib2.DexFileFactory;
import org.jf.dexlib2.Opcodes;
import org.jf.dexlib2.iface.ClassDef;
import org.jf.dexlib2.iface.DexFile;
import org.jf.dexlib2.iface.Field;
import org.jf.dexlib2.iface.Method;
import org.jf.dexlib2.immutable.ImmutableDexFile;
import pxb.android.axml.Util;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.util.*;
import java.util.stream.Collectors;

public class DexMerge {

    // 核心限制：方法引用数 + 字段引用数 + 类总数 阈值（预留缓冲）
    private static final int TOTAL_LIMIT = 65535;

    public void merge(String aApkPath, String bApkPath, String outApkPath, Map<String, String> idsMap) throws Exception {

        // 校验APK路径
        if (!new File(aApkPath).exists() || !aApkPath.endsWith(".apk")) {
            throw new Exception("aApkPath路径错误或不是APK文件");
        }
        List<byte[]> aDexDatas = extractAllDexFromApk(aApkPath);
        if (aDexDatas.isEmpty()) {
            throw new Exception("aApkPath未加载到dex文件");
        }

        if (!new File(bApkPath).exists() || !bApkPath.endsWith(".apk")) {
            throw new Exception("bApkPath路径错误或不是APK文件");
        }
        List<byte[]> bDexDatas = extractAllDexFromApk(bApkPath);
        if (bDexDatas.isEmpty()) {
            throw new Exception("bApkPath未加载到dex文件");
        }

        // 加载并合并类定义
        List<ClassDef> aDexClassList = loadClassDefs(aDexDatas);
        List<ClassDef> bDexClassList = loadClassDefs(bDexDatas);

        DexEditor.modifyFakerAppClass(aDexClassList);

        //ModifiyDexClassID1.modifiyClassIds(bDexClassList, idsMap);
        ModifiyDexClassID1.modifiyClassIds(bDexClassList, idsMap);
        List<ClassDef> allClasses = new ArrayList<>(aDexClassList);
        mergeUniqueClasses(allClasses, aDexClassList, bDexClassList);

        // 排序并写入Dex文件
        sortAndWriteDexFiles(allClasses,outApkPath);
    }

    /**
     * 从APK文件中提取所有dex文件并返回字节数组列表
     * @param apkPath APK文件的绝对路径
     * @return 包含所有dex文件内容的字节数组列表
     * @throws Exception  读取文件内容时发生IO错误
     */
    public static List<byte[]> extractAllDexFromApk(String apkPath) throws Exception {
        List<byte[]> dexList = new ArrayList<>();
        ZipFile zipFile = new ZipFile(apkPath);
        List<FileHeader> fileHeaders = zipFile.getFileHeaders();
        for (FileHeader header : fileHeaders) {
            String entryName = header.getFileName();
            if (entryName != null && entryName.toLowerCase().endsWith(".dex")) {
                InputStream is = zipFile.getInputStream(header);
                dexList.add(is.readAllBytes());
                is.close();
            }
        }
        return dexList;
    }

    /**
     * 从Dex字节数组加载类定义
     */
    private List<ClassDef> loadClassDefs(List<byte[]> bytes) throws IOException {
        List<ClassDef> classDefs = new ArrayList<>();
        for (byte[] bt : bytes) {
            DexFile dexFile = DexFileFactory.loadDexFile(bt, null);
            classDefs.addAll(dexFile.getClasses());
        }
        return classDefs;
    }

    /**
     * 合并唯一类（避免重复）
     */
    private void mergeUniqueClasses(List<ClassDef> targetList, List<ClassDef> aClasses, List<ClassDef> bClasses) {
        Set<String> aClassTypes = aClasses.stream()
                .map(ClassDef::getType)
                .collect(Collectors.toSet());

        for (ClassDef bClass : bClasses) {
            if (!aClassTypes.contains(bClass.getType())) {
                targetList.add(bClass);
            }
        }
    }

//    /**
//     * 排序并按限制拆分写入Dex文件
//     */
//    private void sortAndWriteDexFiles(List<ClassDef> allClasses) throws IOException {
//        // 按包名排序，确保同包类尽量在同一Dex
//        Collections.sort(allClasses, (c1, c2) -> {
//            String pkg1 = DexEditor.getTypePackage(c1.getType());
//            String pkg2 = DexEditor.getTypePackage(c2.getType());
//            return pkg1.compareTo(pkg2);
//        });
//
//        // 按总限制拆分类组
//        List<List<ClassDef>> dexGroups = splitByTotalLimit(allClasses);
//        System.out.println("拆分完成，共生成 " + dexGroups.size() + " 个Dex组");
//
//        // 创建输出目录
//        File outDir = new File(outDirPath);
//        if (!outDir.exists()) {
//            outDir.mkdirs();
//        }
//
//        // 写入Dex文件
//        Opcodes opcodes = Opcodes.getDefault();
//        for (int i = 0; i < dexGroups.size(); i++) {
//            String dexFileName = i == 0 ? "classes.dex" : "classes" + (i + 1) + ".dex";
//            File outFile = new File(outDir, dexFileName);
//            List<ClassDef> group = dexGroups.get(i);
//            System.out.println("写入Dex文件: " + dexFileName + "，包含 " + group.size() + " 个类");
//
//            // 验证组合法性
//            validateGroup(group);
//            // 直接写入（无递归）
//            writeDexFile(outFile, group, opcodes);
//        }
//    }

    /**
     * 排序并按限制拆分写入Dex文件，使用临时文件提高安全性
     */
    private void sortAndWriteDexFiles(List<ClassDef> allClasses, String outApkPath) throws IOException {
        Collections.sort(allClasses, (c1, c2) -> {
            String pkg1 = DexEditor.getTypePackage(c1.getType());
            String pkg2 = DexEditor.getTypePackage(c2.getType());
            return pkg1.compareTo(pkg2);
        });

        List<List<ClassDef>> dexGroups = splitByTotalLimit(allClasses);
        System.out.println("拆分完成，共生成 " + dexGroups.size() + " 个Dex组");
        Opcodes opcodes = Opcodes.getDefault();
        for (int i = 0; i < dexGroups.size(); i++) {
            String dexFileName = i == 0 ? "classes.dex" : "classes" + (i + 1) + ".dex";
            List<ClassDef> group = dexGroups.get(i);
            System.out.println("处理Dex文件: " + dexFileName + "，包含 " + group.size() + " 个类");
            validateGroup(group);
            File tempFile = File.createTempFile("dex-temp-", ".dex");
            try {
                writeDexFile(tempFile, group, opcodes);
                Zip4jUtil.writeFileBytes(outApkPath,dexFileName,Util.readFile(tempFile),false);
            } finally {
                if (tempFile.exists() && !tempFile.delete()) {
                    System.err.println("无法清理临时文件: " + tempFile.getPath());
                }
            }
        }
    }


    /**
     * 按总限制（方法+字段+类数）拆分类列表
     * 核心逻辑：仅根据 方法引用数 + 字段引用数 + 类总数 是否超过65536进行拆分
     */
    private List<List<ClassDef>> splitByTotalLimit(List<ClassDef> allClasses) {
        List<List<ClassDef>> groups = new ArrayList<>();
        List<ClassDef> currentGroup = new ArrayList<>();

        // 用于统计当前组的方法、字段引用数（去重）
        Set<String> methodSet = new HashSet<>();
        Set<String> fieldSet = new HashSet<>();

        for (ClassDef classDef : allClasses) {
            // 计算添加当前类后新增的方法和字段数（去重后）
            int newMethods = countNewMethods(methodSet, classDef);
            int newFields = countNewFields(fieldSet, classDef);

            // 计算添加后总限制：方法引用数 + 字段引用数 + 类总数
            int totalAfterAdd = (methodSet.size() + newMethods) + (fieldSet.size() + newFields) + (currentGroup.size() + 1);

            // 如果超过限制，拆分当前组
            if (totalAfterAdd > TOTAL_LIMIT) {
                groups.add(currentGroup);
                // 初始化新组
                currentGroup = new ArrayList<>();
                methodSet.clear();
                fieldSet.clear();
                // 重新计算当前类的新增数（新组为空）
                newMethods = countNewMethods(methodSet, classDef);
                newFields = countNewFields(fieldSet, classDef);
            }

            // 添加到当前组并更新统计
            currentGroup.add(classDef);
            addMethodsToSet(methodSet, classDef);
            addFieldsToSet(fieldSet, classDef);
        }

        // 添加最后一个组
        if (!currentGroup.isEmpty()) {
            groups.add(currentGroup);
        }

        return groups;
    }

    /**
     * 验证组是否符合总限制
     */
    private void validateGroup(List<ClassDef> group) {
        Set<String> methodSet = new HashSet<>();
        Set<String> fieldSet = new HashSet<>();
        for (ClassDef classDef : group) {
            addMethodsToSet(methodSet, classDef);
            addFieldsToSet(fieldSet, classDef);
        }
        int total = methodSet.size() + fieldSet.size() + group.size();
        System.out.println("组验证 - 方法引用数: " + methodSet.size() + ", 字段引用数: " + fieldSet.size() + ", 类总数: " + group.size() + ", 总和: " + total + " (限制值: " + TOTAL_LIMIT + ")");

        if (total > TOTAL_LIMIT) {
            System.err.println("警告: 组总和超过限制，可能导致Dex生成失败！");
        }
    }

    /**
     * 写入Dex文件（无递归逻辑）
     */
    private void writeDexFile(File outFile, List<ClassDef> classDefs, Opcodes opcodes) throws IOException {
        try {
            ImmutableDexFile dexFile = new ImmutableDexFile(opcodes, classDefs);
            DexFileFactory.writeDexFile(outFile.getAbsolutePath(), dexFile);
        } catch (Exception e) {
            System.err.println("写入Dex文件失败: " + outFile.getName() + "，原因: " + e.getMessage());
            throw e; // 直接抛出异常，不进行递归拆分
        }
    }

    // 以下为方法和字段的统计辅助方法
    private int countNewMethods(Set<String> methodSet, ClassDef classDef) {
        int count = 0;
        for (Method method : classDef.getMethods()) {
            String key = getMethodKey(method);
            if (!methodSet.contains(key)) {
                count++;
            }
        }
        return count;
    }

    private void addMethodsToSet(Set<String> methodSet, ClassDef classDef) {
        for (Method method : classDef.getMethods()) {
            methodSet.add(getMethodKey(method));
        }
    }

    private int countNewFields(Set<String> fieldSet, ClassDef classDef) {
        int count = 0;
        for (Field field : classDef.getStaticFields()) {
            String key = getFieldKey(field);
            if (!fieldSet.contains(key)) {
                count++;
            }
        }
        for (Field field : classDef.getInstanceFields()) {
            String key = getFieldKey(field);
            if (!fieldSet.contains(key)) {
                count++;
            }
        }
        return count;
    }

    private void addFieldsToSet(Set<String> fieldSet, ClassDef classDef) {
        for (Field field : classDef.getStaticFields()) {
            fieldSet.add(getFieldKey(field));
        }
        for (Field field : classDef.getInstanceFields()) {
            fieldSet.add(getFieldKey(field));
        }
    }

    // 生成方法唯一标识（用于去重统计）
    private String getMethodKey(Method method) {
        String params = getParametersAsString(method.getParameterTypes());
        return method.getDefiningClass() + "->" + method.getName() + method.getReturnType() + params;
    }

    // 生成字段唯一标识（用于去重统计）
    private String getFieldKey(Field field) {
        return field.getDefiningClass() + "->" + field.getName() + ":" + field.getType();
    }

    // 将方法参数列表转为字符串
    private String getParametersAsString(List<? extends CharSequence> parameters) {
        StringBuilder sb = new StringBuilder();
        sb.append('(');
        for (int i = 0; i < parameters.size(); i++) {
            if (i > 0) {
                sb.append(',');
            }
            sb.append(parameters.get(i));
        }
        sb.append(')');
        return sb.toString();
    }

    public static void main(String[] args) throws Exception {
        String aApkPath = "C:/Users/15666/Desktop/test/demo2.apk";
        String bApkPath = "C:/Users/15666/Desktop/test/xiaomi.apk";
        String outDirPath = "C:/Users/15666/Desktop/test/merged/";
        DexMerge dexUtil = new DexMerge();
        Map<String, String> ids = new HashMap<>();
        ids.put("0x7f0b001c", "0x7f000001");
        dexUtil.merge(aApkPath, bApkPath, outDirPath, ids);
    }
}